import os
from ament_index_python.packages import get_package_share_directory
from launch import LaunchDescription
from launch.actions import DeclareLaunchArgument, IncludeLaunchDescription
from launch.conditions import IfCondition
from launch.substitutions import LaunchConfiguration, PathJoinSubstitution
from launch.launch_description_sources import PythonLaunchDescriptionSource
from launch_ros.actions import Node


def generate_launch_description():
    # Get package share directories
    pkg_dirs = {
        "go2_description": get_package_share_directory("go2_description"),
        "go2_driver_py": get_package_share_directory("go2_driver_py"),
        "go2_perception": get_package_share_directory("go2_perception"),
        "go2_slam": get_package_share_directory("go2_agiros_slam"),
        "go2_cartographer": get_package_share_directory("go2_agiros_cartographer"),
    }

    # Launch arguments
    use_sim_time_arg = DeclareLaunchArgument(
        'use_sim_time',
        default_value='false',
        description='Use simulation (Gazebo) clock if true'
    )
    
    use_rviz_arg = DeclareLaunchArgument(
        'use_rviz',
        default_value='true',
        description='Whether to launch RViz'
    )

    # Nodes
    cartographer_node = Node(
        package='cartographer_ros',
        executable='cartographer_node',
        parameters=[{'use_sim_time': LaunchConfiguration('use_sim_time')}],
        arguments=[
            '-configuration_directory', 
            PathJoinSubstitution([pkg_dirs["go2_cartographer"], "params"]),
            '-configuration_basename', 'mydog.lua'
        ],
        output='screen'
    )

    cartographer_occupancy_grid_node = Node(
        package='cartographer_ros',
        executable='cartographer_occupancy_grid_node',
        parameters=[
            {'use_sim_time': LaunchConfiguration('use_sim_time')},
            {'resolution': 0.05}
        ],
        output='screen'
    )

    driver_node = Node(
        package="go2_driver_py",
        executable="driver_2",
        parameters=[PathJoinSubstitution([
            pkg_dirs["go2_driver_py"], "params", "go2_driver.yaml"
        ])],
        output="screen"
    )

    rviz_node = Node(
        package="rviz2",
        executable="rviz2",
        arguments=["-d", PathJoinSubstitution([
            pkg_dirs["go2_slam"], "rviz", "go2.rviz"
        ])],
        condition=IfCondition(LaunchConfiguration("use_rviz")),
        output='screen'
    )

    # Launch descriptions
    description_launch = IncludeLaunchDescription(
        PythonLaunchDescriptionSource(
            PathJoinSubstitution([
                pkg_dirs["go2_description"], "launch", "display.launch.py"
            ])
        ),
        launch_arguments={"use_join_state_publisher": "false"}.items()
    )

    perception_launch = IncludeLaunchDescription(
        PythonLaunchDescriptionSource(
            PathJoinSubstitution([
                pkg_dirs["go2_perception"], "launch", "go2_pointcloud_process.launch.py"
            ])
        )
    )

    return LaunchDescription([
        use_sim_time_arg,
        use_rviz_arg,
        
        # 1. Start URDF model and TF tree
        description_launch,
        
        # 2. Start lidar point cloud processing (depends on TF tree)
        perception_launch,
        
        # 3. Start driver node (publishes odometry)
        driver_node,
        
        # 4. Cartographer nodes
        cartographer_node,
        cartographer_occupancy_grid_node,
        
        # 5. RViz visualization (start last)
        rviz_node,
    ])